#include <cassert>
#include <ctime>
#include <iostream>
#include <list>
#include <vector>
#include <queue>
#include <map>
#include <string>
#include "itkImage.h"
#include "graph.h"
#include "alltypedefs.h"
#include "Correspondence.h"
#include "CalcCorrelations.h"

/* #define USE_OMP */
/* #define DEBUG */
#define GREEDY

enum VerboseLevel {
  VERBOSE_INFO = 0,
  VERBOSE_NONE,
  VERBOSE_LOW,
  VERBOSE_MEDIUM,
  VERBOSE_HIGH
};

enum NodeStatus {
  NODE_NORMAL = 0, 
  NODE_DELETED
};

enum TerminalType {
  TERMINAL_NONE = 0,
  TERMINAL_SOURCE,
  TERMINAL_SINK
};

enum GraphType {
  GRAPH_NONE = 0, 
  GRAPH_FWD,
  GRAPH_BWD
};

enum SegmentationType {
  SEG_NONE = 0,
  SEG_FWD, 
  SEG_BWD, 
  SEG_BOTH
};

const double FLOAT_INF = 999999.0;
const unsigned int IMAGE_DIMENSION = 2;
typedef unsigned int IntPixel;
typedef IntPixel LabelPixel;
typedef itk::Image<LabelPixel, IMAGE_DIMENSION> LabelImage;
typedef unsigned char GrayPixel;
typedef itk::Image<GrayPixel, IMAGE_DIMENSION> GrayImage;
typedef itk::RGBPixel<GrayPixel> RgbPixel;
typedef itk::Image<RgbPixel, IMAGE_DIMENSION> RgbImage;
typedef Graph<double, double, double> RawGraph;
typedef std::map< int, std::vector<CorrespondenceMatrix*> > VolumeCorrespondenceSet;
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
class Point
{
 public:
  int x;
  int y;
  Point (int x = -1, int y = -1);
  Point (LabelImage::IndexType index);
  void set (int x, int y);
  bool isEqual (Point point);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::list<Point> ComponentPointSet;
class Component
{
 public:
  int no;
  ComponentPointSet points;
  bool isOnBoundary;
  Component (int no = -1, bool isOnBoundary = true);
  Component (const Component& component);
  void set (int no, bool isOnBoundary);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::vector<Component> LayerComponentSet;
class Layer
{
 public:
  int no;
  LayerComponentSet components;
  LabelImage::Pointer labelImage;
  Layer (int no = -1, LabelImage::Pointer labelImage = NULL);
  Layer (const Layer& layer);
  void set (int no, LabelImage::Pointer labelImage);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::vector<Layer> VolumeLayerSet;
class Volume
{
 public:
  int startSliceNo;
  int startLayerNo;
  int layerNum;
  VolumeLayerSet layers;
  Volume (int startSliceNo = -1, int startLayerNo = -1, int layerNum = -1);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
class Edge
{
 public:
  int endNodeNo;
  int area;			/* Overlapping area with endNodeNo */
  double weight_area;		/* Weight of overlapping */
  double correlation;
  double weight_correlation;	/* Weight of correlation */
  double weight;		/* Total weight */
  Edge (int endNodeNo = -1, int area = -1, double weight_area = -1.0, 
	double correlation = -1.0, double weight_correlation = -1.0, 
	double weight = -1.0);
  void set (int endNodeNo, int area, double weight_area, 
	    double correlation, double weight_correlation, 
	    double weight);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::list<Edge> NodeEdgeSet;
class Node
{
 public:
  int no;
  int layerNo;
  int componentNo;
  NodeStatus status;
  NodeEdgeSet edges;
  NodeEdgeSet reverseEdges;
  NodeEdgeSet skipEdges;
  NodeEdgeSet reverseSkipEdges;
  Node (int no = -1, int layerNo = -1, int componentNo = -1, 
	NodeStatus status = NODE_NORMAL);
  Node (const Node& node);
  void set (int no, int layerNo, int componentNo, NodeStatus status);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
class SegmentationNode
{
 public:
  int layerNo;
  int nodeNo;
  SegmentationType segType;
  SegmentationNode (int layerNo = -1, int nodeNo = -1, 
		    SegmentationType segType = SEG_NONE);
  void set (int layerNo, int nodeNo, 
	    SegmentationType segType);
  bool isEqual(SegmentationNode const& segNode);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::list<SegmentationNode> SegGraphSegNodeSet;
class SegmentationGraph
{
 public:
  SegGraphSegNodeSet segNodeSet;
  SegmentationGraph () {}
  SegmentationGraph (const SegmentationGraph& segGraph);
  void addSegNode (SegmentationNode const& segNode, bool isScan = false);
  void sort ();
  SegmentationGraph getIntersection (SegmentationGraph& segGraph);
  SegmentationGraph getUnion (SegmentationGraph& segGraph);
  SegmentationGraph getDifference (SegmentationGraph& segGraph);
  SegmentationGraph getCertainTypeSegGraph (SegmentationType type);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
class RawNode
{
 public:
  int layerNo;
  int nodeNo;
  RawNode (int layerNo = -1, int nodeNo = -1);
  void set (int layerNo, int nodeNo);
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
typedef std::vector<Node> GraphLayerNodeSet;
typedef std::vector<GraphLayerNodeSet> GraphNodeSet;
typedef std::vector<RawNode> NodeCorrespondence;
typedef std::vector< std::map<int, int> > RawNodeCorrespondence;
class Graph3D
{
 public:
  GraphNodeSet nodes;
  RawGraph* body;
  GraphType graphType;
  int bodyStartLayerNo;		/* Note first layer # for body */
  int bodyLayerNum;		/* Note # of layers for body */
  int bodyNodeNum;
  int bodyEdgeNum;
  /* Map: rawNodeNo --> [layerNo][nodeNo] */
  NodeCorrespondence nodeCorrespondence; 
  /* Map: [layerNo - bodyStartLayerNo][nodeNo] --> rawNodeNo */
  RawNodeCorrespondence rawNodeCorrespondence; 
  Graph3D ();
  ~Graph3D();
  void changeNodeStatus (SegmentationGraph const& segGraph, NodeStatus status);
  void addNode (Node const& node);
  void addEdge (int layerNo, int nodeNo, Edge const& edge);
  void addSkipEdge (int layerNo, int nodeNo, Edge const& edge);
  void create (int bodyStartLayerNo, int bodyLayerNum, GraphType graphType);
  void create (int startLayerNo, int layerNum);
  void destroy ();
};
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
/* util.cxx */
std::string makePath (std::string pathTemplate, int sliceNo0, int sliceNo1 = -1);
LabelImage::Pointer readLabelImage (const char* path);
LabelPixel getLabel (LabelImage::Pointer labelImage, int x, int y);
void writeArrayBinary (int* array, int length, const char* filePath);
void readArrayBinary (int* array, int length, const char* filePath);

/* prep.cxx */
void generateVolume (std::vector<std::string> paths, Volume& volume, 
		     VerboseLevel verbose = VERBOSE_LOW);

/* construct.cxx */
void generateGraph (Volume const& volume, Graph3D& graph, 
		    std::vector<std::string> const& adjacentCorrelationPaths, 
		    std::vector<std::string> const& skipCorrelationPaths, 
		    std::vector<std::string> const& adjacentAreaPaths, 
		    std::vector<std::string> const& skipAreaPaths, 
		    double sigma_area, double sigma_correlation, double alpha, 
		    VerboseLevel verbose = VERBOSE_LOW);

/* segment.cxx */
void segment (Graph3D& graph, std::vector<int> fwdStartComponentNoSet, 
	      int startLayerNo, int layerNum, int pileLayerNum, 
	      SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	      double& fwdMaxFlow, double& bwdMaxFlow, 
	      VerboseLevel verbose = VERBOSE_LOW);
void segment (Graph3D& graph, std::vector<int> fwdStartComponentNoSet, 
	      int startLayerNo, int layerNum, 
	      SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	      double& fwdMaxFlow, double& bwdMaxFlow, 
	      VerboseLevel verbose = VERBOSE_LOW);
void extractGroup (Volume const& volume, Graph3D const& graph, int startLayerNo, 
		   SegmentationGraph fwdSegGraph, SegmentationGraph bwdSegGraph, 
		   std::vector<int> fwdStartComponentNoSet, 
		   SegmentationGraph& resSegGraph);

/* vis.cxx */
void fillComponent (Component const& component, RgbImage::Pointer image, 
		    RgbPixel const& fgPixel, double bgPortion);
RgbImage::Pointer GrayToRgb (const char* imageFilePath);
void writeRgbImage (RgbImage::Pointer rgbImage, const char* imageFilePath);
void drawRGB (Volume const& volume, Graph3D const& graph, 
	      SegmentationGraph fwdSegGraph, 
	      SegmentationGraph bwdSegGraph, 
	      std::vector<std::string> backgroundPaths, 
	      std::vector<std::string> outputPaths, 
	      VerboseLevel verbose = VERBOSE_LOW);

/* operation.cxx */
void operation (int argc, char* argv[]);

/* xml.cxx */
void computeComponentInsideCenter (LabelImage::Pointer labelImage, 
				   Component const& component, int& xc, int& yc, 
				   VerboseLevel verbose = VERBOSE_LOW);
void writeXML (Volume const& volume, Graph3D const& graph, 
	       SegmentationGraph const& resSegGraph,
	       int startSliceNo, const char* xmlPath, 
	       VerboseLevel verbose = VERBOSE_LOW, int pathID = 0);
